# 预测模块
import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
import sys
sys.path.append('..')
from config import PREDICTION_DAYS

class Predictor:
    def __init__(self):
        # 预测天数配置
        self.prediction_days = PREDICTION_DAYS
        
    def prepare_features(self, df):
        """准备预测特征"""
        # 确保数据按日期排序
        df = df.sort_values('trade_date')
        
        # 创建特征DataFrame
        features = pd.DataFrame()
        
        # 添加技术指标作为特征
        features['close'] = df['close']
        features['open'] = df['open']
        features['high'] = df['high']
        features['low'] = df['low']
        features['vol'] = df['vol']
        
        # 添加移动平均线
        features['ma5'] = df['ma5']
        features['ma10'] = df['ma10']
        features['ma20'] = df['ma20']
        
        # 添加MACD指标
        features['macd'] = df['macd']
        features['macd_signal'] = df['macd_signal']
        features['macd_hist'] = df['macd_hist']
        
        # 添加KDJ指标
        features['k'] = df['k']
        features['d'] = df['d']
        features['j'] = df['j']
        
        # 添加RSI指标
        features['rsi'] = df['rsi']
        
        # 添加布林带指标
        features['boll_upper'] = df['boll_upper']
        features['boll_middle'] = df['boll_middle']
        features['boll_lower'] = df['boll_lower']
        
        # 添加价格变化率
        features['price_change'] = df['close'].pct_change()
        features['price_change_5d'] = df['close'].pct_change(5)
        features['price_change_10d'] = df['close'].pct_change(10)
        
        # 添加成交量变化率
        features['volume_change'] = df['vol'].pct_change()
        features['volume_change_5d'] = df['vol'].pct_change(5)
        
        # 删除NaN值
        features = features.dropna()
        
        return features
    
    def predict_price(self, df, days=1):
        """预测未来价格"""
        if len(df) < 30:  # 确保有足够的历史数据
            return {'error': '历史数据不足，无法进行预测'}
        
        # 准备特征
        features = self.prepare_features(df)
        
        # 准备训练数据
        X = features.iloc[:-days].values  # 使用除最后days天外的所有数据作为特征
        y = features['close'].shift(-days).iloc[:-days].values  # 预测未来days天的收盘价
        
        # 标准化特征
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)
        
        # 训练线性回归模型
        model = LinearRegression()
        model.fit(X_scaled, y)
        
        # 预测
        X_pred = features.iloc[-1:].values
        X_pred_scaled = scaler.transform(X_pred)
        predicted_price = model.predict(X_pred_scaled)[0]
        
        # 计算预测涨跌幅
        current_price = df['close'].iloc[-1]
        change_percent = (predicted_price - current_price) / current_price * 100
        
        return {
            'current_price': current_price,
            'predicted_price': predicted_price,
            'change_percent': change_percent,
            'days': days
        }
    
    def predict_all(self, df):
        """预测所有配置的天数"""
        result = {}
        
        for days in self.prediction_days:
            prediction = self.predict_price(df, days)
            result[f'{days}day_prediction'] = prediction
        
        return result
    
    def calculate_support_resistance(self, df):
        """计算支撑位和阻力位"""
        if len(df) < 20:
            return {'error': '历史数据不足，无法计算支撑阻力位'}
        
        # 获取最近20天的数据
        recent_data = df.iloc[-20:]
        
        # 计算支撑位（取最近20天的低点中的较高值）
        lows = recent_data['low'].sort_values()
        supports = lows.iloc[:3].tolist()  # 取最低的3个值作为支撑位
        
        # 计算阻力位（取最近20天的高点中的较低值）
        highs = recent_data['high'].sort_values(ascending=False)
        resistances = highs.iloc[:3].tolist()  # 取最高的3个值作为阻力位
        
        # 获取当前价格
        current_price = df['close'].iloc[-1]
        
        # 筛选出当前价格以下的支撑位和当前价格以上的阻力位
        valid_supports = [s for s in supports if s < current_price]
        valid_resistances = [r for r in resistances if r > current_price]
        
        return {
            'current_price': current_price,
            'support_levels': valid_supports,
            'resistance_levels': valid_resistances
        }
    
    def generate_trading_signal(self, df):
        """生成交易信号"""
        if len(df) < 2:
            return {'signal': '数据不足，无法生成交易信号'}
        
        # 获取最新数据
        latest_data = df.iloc[-1]
        prev_data = df.iloc[-2]
        
        # 初始化信号计数器
        buy_signals = 0
        sell_signals = 0
        hold_signals = 0
        
        # MACD信号
        if prev_data['macd_hist'] < 0 and latest_data['macd_hist'] > 0:
            buy_signals += 1
        elif prev_data['macd_hist'] > 0 and latest_data['macd_hist'] < 0:
            sell_signals += 1
        else:
            hold_signals += 1
        
        # KDJ信号
        if latest_data['k'] > latest_data['d']:
            if latest_data['k'] < 20:
                buy_signals += 2  # 超卖区域金叉，强烈买入信号
            else:
                buy_signals += 1
        else:
            if latest_data['k'] > 80:
                sell_signals += 2  # 超买区域死叉，强烈卖出信号
            else:
                sell_signals += 1
        
        # RSI信号
        if latest_data['rsi'] < 30:
            buy_signals += 1
        elif latest_data['rsi'] > 70:
            sell_signals += 1
        else:
            hold_signals += 1
        
        # 布林带信号
        if latest_data['close'] < latest_data['boll_lower']:
            buy_signals += 1
        elif latest_data['close'] > latest_data['boll_upper']:
            sell_signals += 1
        else:
            hold_signals += 1
        
        # 移动平均线信号
        if latest_data['ma5'] > latest_data['ma10'] > latest_data['ma20']:
            buy_signals += 1  # 多头排列
        elif latest_data['ma5'] < latest_data['ma10'] < latest_data['ma20']:
            sell_signals += 1  # 空头排列
        else:
            hold_signals += 1
        
        # 根据信号数量生成最终建议
        if buy_signals > sell_signals and buy_signals > hold_signals:
            signal = '买入'
            strength = min(buy_signals / (buy_signals + sell_signals + hold_signals) * 10, 10)
            reason = '技术指标多数显示买入信号'
        elif sell_signals > buy_signals and sell_signals > hold_signals:
            signal = '卖出'
            strength = min(sell_signals / (buy_signals + sell_signals + hold_signals) * 10, 10)
            reason = '技术指标多数显示卖出信号'
        else:
            signal = '持有'
            strength = min(hold_signals / (buy_signals + sell_signals + hold_signals) * 10, 10)
            reason = '技术指标多数显示观望信号'
        
        return {
            'signal': signal,
            'strength': strength,
            'reason': reason,
            'buy_signals': buy_signals,
            'sell_signals': sell_signals,
            'hold_signals': hold_signals
        }